"""
@Time    : 2019/10/15 9:07
@Author  : CcH
"""

from DataLoader.Picture_Process import file_2_allpicture
import torch
from tqdm import tqdm
import os
import pandas as pd
from model.model import LipNet,LipSeqLoss


def read_data(path):
    """key,value"""
    data = {}
    with open(path, 'r', encoding="utf-8") as fr:
        for line in tqdm(fr.readlines()):
            id_text = line.strip().split('\t')
            data[id_text[0]] = id_text[1]
    return data

def remove_len0(path):
    picture_len0 = []
    # path = r"D:\XW_Bank\LipRecognition\train\mouth_train"
    dirlist = os.listdir(path)
    for id in dirlist:
        newpath = path + "//" + id
        length = len(os.listdir(newpath))
        if length != 0:
            picture_len0.append(id)
    return picture_len0




def test(test_path,model_path):
    chinese_info = torch.load("./DataLoader/train_labels.pt")
    word2idx = chinese_info["word2idx"]
    idx2word = {}
    for key, value in word2idx.items():
        idx2word[value] = key
    model = LipNet(313).cuda()
    model.load_state_dict(torch.load(model_path))
    model.eval()

    dirlist = remove_len0(test_path)
    ids, words = [], []
    # criterion = nn.NLLLoss()
    for id in tqdm(dirlist, desc='  - (Training)  '):
        src,pct_len = file_2_allpicture(test_path + "//" + id)
        pct_len = torch.LongTensor([pct_len])
        src = src.cuda()
        pct_len = pct_len.cuda()
        with torch.no_grad():
            output = model(src)
        ids.append(id)

        # average_volumn = output[0, :pct_len[0]].sum(0)
        # _, max_index = torch.max(average_volumn, 0)
        # pred = max_index.data.item()
        pred = output.cpu().max(1)[1].item()
        words.append(idx2word[pred])

    reseult = {"id": ids, "word": words}
    reseult = pd.DataFrame(reseult, columns=['id', 'word'], index=None)
    reseult[["id", "word"]].to_csv("预测结果.csv", index=None,header=0)



if __name__ == '__main__':
    test_path = r"D:\XW_Bank\LipRecognition\test\lip_100_50_test"
    model_path = r"./weight/Lip_Recognition_val_loss.pkl"
    test(test_path,model_path)